from astropy.time import Time
from astropy import units as u
import numpy as np
import json
import os

from pysagin.spacebase import Constellation, Satellite
from pysagin.groundbase import GroundStation
from pysagin.airbase import UavBase, Trajectory
from pysagin.device import RadioFixedAccess, RadioRandomAccess, Radio

# region 初始化全局参数
print("INFO: 初始化全局参数")

start_time = Time("2025-03-27 00:00:00")

# endregion

# region 创建 Walker 星座
print("INFO: 创建 Walker 星座")

constellation = Constellation(
    name="S",
    type="Walker",
    T=66,  # 总卫星数量
    P=6,   # 轨道面数量
    F=0.5, # 相位因子
    a=6371+550,  # 半长轴（单位：km）
    e=0,  # 偏心率
    inc=53,  # 倾角（单位：度）
    init_time=start_time
)

constellation.update_all(start_time)
sats_gps = constellation.get_all_current_gps()
sats_lon_list = [sat_gps[0] for sat_gps in sats_gps]
sats_lat_list = [sat_gps[1] for sat_gps in sats_gps]
sats_alt_list = [sat_gps[2] for sat_gps in sats_gps]

# endregion

# region 创建地面站
print("INFO: 创建地面站")

station_1 = GroundStation(
    name="G1",
    init_time=start_time,
    init_gps=[126.662892, 45.742140, 0.146]
)

station_2 = GroundStation(
    name="G2",
    init_time=start_time,
    init_gps=[122.116342, 37.510357, 0.007]
)

station_3 = GroundStation(
    name="G3",
    init_time=start_time,
    init_gps=[114.057865, 22.543096, 0.004]
)

stations_gps = [station_1.current_gps, station_2.current_gps, station_3.current_gps]
stations_lon_list = [station_gps[0] for station_gps in stations_gps]
stations_lat_list = [station_gps[1] for station_gps in stations_gps]
stations_alt_list = [station_gps[2] for station_gps in stations_gps]

# endregion

# region 创建无人机
print("INFO: 创建无人机")

uav_1_trajectory = Trajectory(name="U1Tr",avg_speed=(100*u.km/u.h).to(u.m/u.s).value)
uav_1_trajectory.add_straight(start=station_1.init_gps, end=station_2.init_gps)
uav_1_trajectory.add_figure8(center=station_2.init_gps, size=2000, loops=10)
uav_1_trajectory.add_straight(start=station_2.init_gps, end=station_1.init_gps)

uav_1 = UavBase(
    name="U1",
    init_time=start_time,
    init_gps=station_1.init_gps
)

uav_1.set_trajectory(trajectory=uav_1_trajectory, 
                     departure_time=start_time + 60 * u.min)

uav_2_trajectory = Trajectory(name="U2Tr",avg_speed=(150*u.km/u.h).to(u.m/u.s).value)
uav_2_trajectory.add_straight(start=station_1.init_gps, end=station_3.init_gps)
uav_2_trajectory.add_circle(center=station_3.init_gps, radius=1000, loops=5)
uav_2_trajectory.add_straight(start=station_3.init_gps, end=station_2.init_gps)

uav_2 = UavBase(
    name="U2",
    init_time=start_time,
    init_gps=station_1.init_gps
)

uav_2.set_trajectory(trajectory=uav_2_trajectory, 
                     departure_time=start_time + 120 * u.min)

uavs_gps = [uav_1.current_gps, uav_2.current_gps]
uavs_lon_list = [uav_gps[0] for uav_gps in uavs_gps]
uavs_lat_list = [uav_gps[1] for uav_gps in uavs_gps]
uavs_alt_list = [uav_gps[2] for uav_gps in uavs_gps]

# endregion

# region 创建无线电台
print("INFO: 创建无线电台")

constellation.create_inter_satellite_links(Device_Radio=RadioFixedAccess, range=10000)


constellation_radio_random = []
for sat in constellation.satellites:
    sat_radio_random = RadioRandomAccess(name=f"{sat.name}D1", platform=sat, range=1000,
                                          max_access_num=-1)
    sat.devices.append(sat_radio_random)
    constellation_radio_random.append(sat_radio_random)


sation_radio_1 = RadioRandomAccess(name=f"{station_1.name}D0", platform=station_1, 
                                   range=1000, max_access_num=-1)
station_1.devices.append(sation_radio_1)
sation_radio_2 = RadioRandomAccess(name=f"{station_2.name}D0", platform=station_2, 
                                   range=1000, max_access_num=-1)
station_2.devices.append(sation_radio_2)
sation_radio_3 = RadioRandomAccess(name=f"{station_3.name}D0", platform=station_3, 
                                   range=1000, max_access_num=-1)
station_3.devices.append(sation_radio_3)


uav_radio_1 = RadioRandomAccess(name=f"{uav_1.name}D0", platform=uav_1, 
                                range=1000, max_access_num=-1)
uav_1.devices.append(uav_radio_1)
uav_radio_2 = RadioRandomAccess(name=f"{uav_2.name}D0", platform=uav_2, 
                                range=1000, max_access_num=-1)
uav_2.devices.append(uav_radio_2)

for sat_radio_random in constellation_radio_random:
    sat_radio_random.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3]+
                                           [uav_radio_1, uav_radio_2])
sation_radio_1.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
sation_radio_2.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
sation_radio_3.add_pre_access_radios([uav_radio_1, uav_radio_2] + constellation_radio_random)
uav_radio_1.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)
uav_radio_2.add_pre_access_radios([sation_radio_1, sation_radio_2, sation_radio_3] + constellation_radio_random)

ALL_DEVICES:list[Radio] = []
for sat in constellation.satellites:
    ALL_DEVICES.extend(sat.devices)
ALL_DEVICES.extend([sation_radio_1, sation_radio_2, sation_radio_3, uav_radio_1, uav_radio_2])

# endregion

# region 导出json数据
print("INFO: 导出json数据")

def monitor_access(current_time: Time):
    constellation.update_all(current_time)
    uav_1.update(current_time)
    uav_2.update(current_time)
    
    data_list = []
    for dev_send in ALL_DEVICES:
        base_send = dev_send.platform
        radio_recvs:list[Radio] = dev_send.get_current_access_radios()
        base_recvs = [radio_recv.platform for radio_recv in radio_recvs]
        if base_recvs == []:
            data = {'node1': base_send.name, 'node2': None, 
                    'range': None}
            data_list.append(data)
        else:
            for base_recv in base_recvs:
                distance = base_send.calculate_distance(base_recv.current_gps, 
                                                        base_send.current_gps)
                data = {'node1': base_send.name, 'node2': base_recv.name, 
                        'range': distance}
                data_list.append(data)
    return data_list

duration = 24 * u.hour  # 总时长
interval = 10 * u.min  # 时间间隔
time_steps = np.arange(0, duration.to(u.s).value, interval.to(u.s).value) * u.s
time_series = start_time + time_steps

if not os.path.exists('data'): os.makedirs('data')

def get_access_data(current_time: Time):
    data_list = monitor_access(current_time)
    data_list = json.dumps(data_list,indent=4)
    
    file_path = f'data/sagin_data_{current_time.to_datetime().strftime("%Y_%m_%d_%H_%M_%S")}'+'.json'
    with open(file_path, 'w') as f:
        f.write(data_list)

from joblib import Parallel, delayed
Parallel(n_jobs=-1)(
    delayed(get_access_data)(current_time) 
    for current_time in time_series
)

# endregion 
